from tonglun_py import *
import numpy as np
import psutil
import gc

# Same function as the original test
export_fn("testpy", lambda x, y, z: x + y + z)

process = psutil.Process()
print('Initial memory:', process.memory_info().rss)

# Test with three arguments like the original
for i in range(1000):
    x = np.array([1.0, 2.0, 3.0])
    result = call_fn("testpy", x, x.copy(), x.copy())
    if i % 100 == 0:
        gc.collect()  # Force garbage collection
        print(f"Iteration {i}, memory: {process.memory_info().rss}")

print('Final memory:', process.memory_info().rss) 